#!/usr/bin/env python3
"""
数据库迁移脚本：添加文件版本管理功能
"""

import sqlite3
import os
from datetime import datetime

def migrate_database():
    """执行数据库迁移"""
    db_path = "filemanager.db"
    
    if not os.path.exists(db_path):
        print("数据库文件不存在，请先启动应用程序创建数据库")
        return False
    
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    try:
        print("开始数据库迁移：添加文件版本管理功能")
        
        # 1. 创建文件版本表
        print("1. 创建file_versions表...")
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS file_versions (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                file_id INTEGER NOT NULL,
                version_number INTEGER NOT NULL,
                path VARCHAR(500) NOT NULL,
                size INTEGER NOT NULL,
                mime_type VARCHAR(100),
                checksum VARCHAR(32) NOT NULL,
                comment TEXT,
                created_by INTEGER NOT NULL,
                is_current BOOLEAN DEFAULT FALSE,
                is_deleted BOOLEAN DEFAULT FALSE,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (file_id) REFERENCES files (id),
                FOREIGN KEY (created_by) REFERENCES users (id),
                UNIQUE(file_id, version_number)
            )
        """)
        print("   - file_versions表创建成功")
        
        # 2. 为现有文件创建初始版本
        print("2. 为现有文件创建初始版本...")
        
        # 获取所有现有的文件（非文件夹）
        cursor.execute("""
            SELECT id, path, size, mime_type, owner_id, created_at
            FROM files 
            WHERE file_type = 'file' AND is_deleted = FALSE
        """)
        
        existing_files = cursor.fetchall()
        
        for file_data in existing_files:
            file_id, path, size, mime_type, owner_id, created_at = file_data
            
            # 检查是否已经有版本记录
            cursor.execute("""
                SELECT COUNT(*) FROM file_versions WHERE file_id = ?
            """, (file_id,))
            
            if cursor.fetchone()[0] == 0:
                # 计算文件的MD5校验和
                import hashlib
                checksum = "unknown"
                if os.path.exists(path):
                    try:
                        with open(path, 'rb') as f:
                            checksum = hashlib.md5(f.read()).hexdigest()
                    except Exception as e:
                        print(f"   - 警告：无法计算文件 {path} 的校验和: {e}")
                
                # 创建初始版本
                cursor.execute("""
                    INSERT INTO file_versions 
                    (file_id, version_number, path, size, mime_type, checksum, comment, created_by, is_current, created_at)
                    VALUES (?, 1, ?, ?, ?, ?, '初始版本', ?, TRUE, ?)
                """, (file_id, path, size, mime_type, checksum, owner_id, created_at))
        
        print(f"   - 为 {len(existing_files)} 个文件创建了初始版本")
        
        # 3. 创建索引以提高查询性能
        print("3. 创建索引...")
        cursor.execute("""
            CREATE INDEX IF NOT EXISTS idx_file_versions_file_id 
            ON file_versions(file_id)
        """)
        cursor.execute("""
            CREATE INDEX IF NOT EXISTS idx_file_versions_current 
            ON file_versions(file_id, is_current)
        """)
        cursor.execute("""
            CREATE INDEX IF NOT EXISTS idx_file_versions_version 
            ON file_versions(file_id, version_number)
        """)
        print("   - 索引创建成功")
        
        conn.commit()
        print("\n=== 迁移完成 ===")
        print("文件版本管理功能已成功添加到数据库")
        return True
        
    except Exception as e:
        print(f"迁移过程中发生错误: {e}")
        conn.rollback()
        return False
    finally:
        conn.close()

if __name__ == "__main__":
    success = migrate_database()
    if success:
        print("\n=== 迁移成功 ===")
        print("现在可以使用文件版本管理功能了")
    else:
        print("\n=== 迁移失败 ===")
        print("请检查数据库连接和权限")